import cv2
import numpy as np
from queue import PriorityQueue


# 图像地图的节点类定义
class Node(object):
    def __init__(self, x, y, asked, cost, cache, queuing, father_node, dp, dq):
        self.x = x
        self.y = y
        self.dp = dp  # 节点的dp缓存
        self.dq = dq  # 节点的dq缓存
        self.asked = asked  # 是否被加入过队列
        self.cost = cost  # 到种子结点的代价
        self.cache = cache  # Z和G的加权缓存
        self.queuing = queuing  # 是否在队列中
        self.father_node = father_node  # 该结点的指针指向（已弃用）
        self.fatherNode = None  # 节点指向的父节点

    # 比较方法重写，用于优先队列排序
    def __eq__(self, other):
        return self.cost == other.cost

    def __lt__(self, other):
        return self.cost < other.cost


# 将0-255的Canny灰阶图转为0-1，获取Fz
def fz(row, col):
    global img_canny
    if img_canny[row, col] == 255:
        return 0
    else:
        return 1


# 将获取Fg的值转为函数处理
def fg(row, col):
    global grad_min, grad_max, G
    return 1 - ((G[row, col] - grad_min) / (grad_max - grad_min))


# 辅助Fd计算的函数
def _d_(p):
    global sx, sy
    return np.array([sx[p[0], p[1]], -sy[p[0], p[1]]])


def _l(p, q):
    if np.dot(_d_(p), q - p) >= 0:
        return q - p
    else:
        return p - q


def _dp(p, q):
    # print("result of _d_ is" + str(_d_(p)))
    # print("result of _l is" + str(_l(p, q)))
    return np.dot(_d_(p), _l(p, q))


def _dq(p, q):
    return np.dot(_l(p, q), _d_(p))


# 计算Fd（必须在计算完maxDp/maxDq之后调用）
def fd(p, q):
    # print("result of dp is" + str(_dp(p, q)))
    global maxDp, maxDq
    dp = img_nodeMap[q[0]][q[1]].dp
    dq = img_nodeMap[q[0]][q[1]].dq
    if maxDp != 0:
        dp /= maxDp
    if maxDq != 0:
        dq /= maxDq
    # if dp != 0 or dq != 0:
    # print("point" + str(q) + ",dp=" + str(dp) + ",dq=" + str(dq))
    return (1 / np.pi) * (np.arccos(dp) + np.arccos(dq))


# 计算整张图的Fd并最终得出cost
def calcGlobalCost():
    global img_nodeMap, seedX, seedY, maxDp, maxDq
    maxDp = 0
    maxDq = 0
    # 计算整张图的初始Dp和Dq
    print("Working on pre-calc...")
    for m in range(img_origin.shape[0]):
        for n in range(img_origin.shape[1]):
            img_nodeMap[m][n].dp = _dp(np.array([seedX, seedY]), np.array([m, n]))
            img_nodeMap[m][n].dq = _dq(np.array([seedX, seedY]), np.array([m, n]))
            # 分别获取整张图的Dp Dq最大值和最小值，用于后续Fd的计算
            if maxDp < np.abs(img_nodeMap[m][n].dp):
                maxDp = np.abs(img_nodeMap[m][n].dp)
            if maxDq < np.abs(img_nodeMap[m][n].dq):
                maxDq = np.abs(img_nodeMap[m][n].dq)
        # print("Pre-calc progress:"+str(m/img_origin.shape[0]))
    # print("maxDp=" + str(maxDp) + ",maxDq=" + str(maxDq))
    print("Pre-calc Done...")
    # 计算最终cost
    for m in range(img_origin.shape[0]):
        for n in range(img_origin.shape[1]):
            img_nodeMap[m][n].cost = img_nodeMap[m][n].cache + 0.43 * fd(np.array([seedX, seedY]), np.array([m, n]))
    # seed的cost清零
    img_nodeMap[seedX][seedY].cost = 0
    print("Cost calc Done...")


# 输入箭头方向和当前坐标获取父节点坐标（已废弃）
def getFatherX(x, arrow):
    if arrow < -8:
        return x - 1
    elif arrow > 8:
        return x + 1
    else:
        return x


def getFatherY(y, arrow):
    if arrow == -11 or arrow == -1 or arrow == 9:
        return y - 1
    elif arrow == -9 or arrow == 1 or arrow == 11:
        return y + 1
    else:
        return y


# 为地图计算最短路径
def getBestWays():
    global img_nodeMap, img_origin, img_draw
    # 统计未被访问的节点数（已废弃）
    unAskedNodes = img_origin.shape[0] * img_origin.shape[1] - 1
    # 创建并初始化Dij队列并将seed推入
    que = PriorityQueue(maxsize=0)
    que.put(img_nodeMap[seedX][seedY])
    print("Running dijkstra...")
    while not que.empty():
        # print(unAskedNodes)
        # 取出队列中cost最小的Node进行访问
        thisNode = que.get()
        # img_draw[thisNode.x, thisNode.y] = np.array([255, 0, 0])
        img_nodeMap[thisNode.x][thisNode.y].queuing = False
        img_nodeMap[thisNode.x][thisNode.y].asked = True
        # 依次访问此Node的周围8个Node
        for i in range(3):
            for j in range(3):
                # 计算目标Node的实际坐标
                ax = i + thisNode.x - 1
                ay = j + thisNode.y - 1
                # 排除超范围的Node
                if ax <= 0 or ay <= 0 or ax >= img_origin.shape[0] - 1 or ay >= img_origin.shape[1] - 1:
                    continue
                # 排除已访问或是在队列中的Node
                if img_nodeMap[ax][ay].asked or img_nodeMap[ax][ay].queuing:
                    continue
                # 对Node进行访问
                unAskedNodes -= 1
                # 判断中心Node是否为最佳通路
                if img_nodeMap[ax][ay].father_node == 0 or img_nodeMap[ax][ay].fatherNode.cost > thisNode.cost:
                    img_nodeMap[ax][ay].fatherNode = thisNode
                    img_nodeMap[ax][ay].father_node = (ax - thisNode.x) + (ay - thisNode.y) * 10
                # 如果Node不在队列中则将其加入队列
                if not img_nodeMap[ax][ay].queuing:
                    img_nodeMap[ax][ay].queuing = True
                    que.put(img_nodeMap[ax][ay])
                '''if img_nodeMap[ax][ay].father_node == 0:
                    print("alert:thisNode.x=" + str(thisNode.x) + ",y=" + str(thisNode.y) + ",askingNode.x=" + str(
                            ax) + ",y=" + str(y))
                    return'''
    print("Dijkstra Done...")
    # 展示Dijkstra的查找结果
    '''img_draw = np.zeros(img_origin.shape, np.uint8)
    for m in range(img_origin.shape[0]):
        for n in range(img_origin.shape[1]):
            img_draw[m, n] = np.array([(getFatherX(m, img_nodeMap[m][n].father_node) - m + 1) * 85,
                                       (getFatherY(n, img_nodeMap[m][n].father_node) - n + 1) * 85, 0])
    cv2.imshow("wind",img_draw)
    cv2.waitKey(0)'''
    return


# 鼠标回调函数
def onMouse(event, x, y, flags, param):
    global seedX, seedY, seedX_origin, seedY_origin, img_draw, img_origin, img_nodeMap, processing, img_draw_former
    # 如果上一次的seed点定位任务计算尚未完成则拒绝工作
    if processing:
        return
    # 处理左键单击（设置seed）
    if event == cv2.EVENT_LBUTTONDOWN:
        if seedX == -1:
            # 初次设置seed点
            processing = True
            seedX = y
            seedY = x
            seedX_origin = seedX
            seedY_origin = seedY
            print("Set origin seed:" + str(seedX) + "," + str(seedY))
            calcGlobalCost()
            '''
            # 打印区域采样数据
            for m in range(5):
                for n in range(5):
                    ax=seedX+m-2
                    ay=seedY+n-2
                    print(img_nodeMap[ax][ay].cost)
            # 图像化展示cost
            for m in range(img_origin.shape[0]):
                for n in range(img_origin.shape[1]):
                    cst = img_nodeMap[m][n].cost * 255
                    img_draw[m, n] = np.array([cst, cst, cst])
            cv2.imshow("wind", img_draw)
            cv2.waitKey(0)
            print(img_nodeMap[seedX-1][seedY-1].cost)
            print(img_nodeMap[seedX-1][seedY].cost)
            print(img_nodeMap[seedX-1][seedY+1].cost)
            print(img_nodeMap[seedX][seedY-1].cost)
            print(img_nodeMap[seedX][seedY].cost)
            print(img_nodeMap[seedX][seedY+1].cost)
            print(img_nodeMap[seedX+1][seedY-1].cost)
            print(img_nodeMap[seedX+1][seedY].cost)
            print(img_nodeMap[seedX+1][seedY+1].cost)
            '''
            getBestWays()
            processing = False
            print("Finish")
        else:
            # 非首次设置seed
            # 将上次的寻路图像作为新的画布
            img_draw_former = img_draw.copy()
            # 清空地图上部分需要重新计算的值
            for m in range(img_origin.shape[0]):
                for n in range(img_origin.shape[1]):
                    img_nodeMap[m][n].cost = -1
                    img_nodeMap[m][n].fatherNode = None
                    img_nodeMap[m][n].father_node = 0
                    img_nodeMap[m][n].asked = False
                    img_nodeMap[m][n].queuing = False
            processing = True
            # 获取鼠标所在位置（需要翻转）
            seedX = y
            seedY = x
            print("Set step seed:" + str(seedX) + "," + str(seedY))
            calcGlobalCost()
            getBestWays()
            processing = False
            print("Finish")
        return

    if event == cv2.EVENT_MOUSEMOVE:
        # 处理鼠标移动
        # 未设置seed时拒绝工作
        if seedX == -1:
            return
        # 获取鼠标位置（需要翻转）
        m = y
        n = x
        # print("Get Cursor:x=" + str(m) + ",y=" + str(n) + ",Seed:x=" + str(seedX) + ",y=" + str(seedY))
        # 刷新画布（清理上次的绘制痕迹）
        img_draw = img_draw_former.copy()
        # 为了防止seed的指针不为0造成死循环和报错，重新设置一遍
        img_nodeMap[seedX][seedY].father_node = 0
        # 绘制导向seed的路径
        while img_nodeMap[m][n].father_node != 0:
            # 画笔颜色为绿色
            img_draw[m, n] = np.array([0, 255, 0])
            '''m = getFatherX(m, img_nodeMap[m][n].father_node)
            n = getFatherY(n, img_nodeMap[m][n].father_node)'''
            # 获取父像素的坐标（需要缓存）
            m_tmp = img_nodeMap[m][n].fatherNode.x
            n_tmp = img_nodeMap[m][n].fatherNode.y
            m = m_tmp
            n = n_tmp
            # 查看是否已到达seed，如果到达则直接break
            if m == seedX and n == seedY:
                # print("Reach seed")
                img_draw[m, n] = np.array([0, 255, 0])
                break
            # print("Turn to " + str(m) + "," + str(n))

            # cv2.waitKey(0)

    if event == cv2.EVENT_RBUTTONDOWN:
        # 处理双击左键，结束工作
        # 未设置seed时拒绝工作
        if seedX == -1:
            return
        # 记忆上次的寻路图像
        img_draw_former = img_draw.copy()
        # 清空地图上需要重新计算的值
        for m in range(img_origin.shape[0]):
            for n in range(img_origin.shape[1]):
                img_nodeMap[m][n].cost = -1
                img_nodeMap[m][n].fatherNode = None
                img_nodeMap[m][n].father_node = 0
                img_nodeMap[m][n].asked = False
                img_nodeMap[m][n].queuing = False
        processing = True
        # 设置初始seed为当前seed并计算
        seedX = seedX_origin
        seedY = seedY_origin
        calcGlobalCost()
        getBestWays()
        # 将当前鼠标位置传入为q寻路（需要翻转）
        m = y
        n = x
        img_draw = img_draw_former.copy()
        img_nodeMap[seedX][seedY].father_node = 0
        while img_nodeMap[m][n].father_node != 0:
            img_draw[m, n] = np.array([0, 255, 0])
            '''m = getFatherX(m, img_nodeMap[m][n].father_node)
            n = getFatherY(n, img_nodeMap[m][n].father_node)'''
            m_tmp = img_nodeMap[m][n].fatherNode.x
            n_tmp = img_nodeMap[m][n].fatherNode.y
            m = m_tmp
            n = n_tmp
            if m == seedX and n == seedY:
                # print("Reach seed")
                img_draw[m, n] = np.array([0, 255, 0])
                break
        print("Finish task")
    return


# 建立图片展示窗口
cv2.namedWindow("wind")
# 读取图片
img_origin = cv2.imread("7.jpg")
print("Initializing...")
# 高斯滤波降噪
img_gaus = cv2.GaussianBlur(img_origin, (1, 1), 2)
# 转换为灰阶图像
img_gray = cv2.cvtColor(img_gaus, cv2.COLOR_BGR2GRAY)
# 使用Canny算子绘制轮廓
img_canny = cv2.Canny(img_gaus, 50, 150)
# 展示Canny算子计算结果
# cv2.imshow("wind", img_canny)
# cv2.waitKey(0)
# 计算图像梯度
sx = cv2.convertScaleAbs(cv2.Sobel(img_gray, cv2.CV_64F, 1, 0, ksize=3))
sy = cv2.convertScaleAbs(cv2.Sobel(img_gray, cv2.CV_64F, 0, 1, ksize=3))
G = cv2.addWeighted(sx, 0.5, sy, 0.5, 0)
# 获取最小、大梯度值
grad_min = G[0, 0]
grad_max = G[0, 0]
for row in range(0, img_origin.shape[0]):
    for col in range(0, img_origin.shape[1]):
        if grad_min > G[row, col]:
            grad_min = G[row, col]
        if grad_max < G[row, col]:
            grad_max = G[row, col]
# 展示梯度计算结果
# cv2.imshow("wind", G)
# cv2.waitKey(0)
# 新建空白节点地图
img_nodeMap = [[Node(m, n, False, -1, 0, False, 0, -1, -1) for n in range(img_origin.shape[1])] for m in
               range(img_origin.shape[0])]
# 计算Z和G的加权值存入缓存
for m in range(img_origin.shape[0]):
    for n in range(img_origin.shape[1]):
        img_nodeMap[m][n].cache = 0.43 * fz(m, n) + 0.14 * fg(m, n)
# 设置鼠标事件回调函数
cv2.setMouseCallback("wind", onMouse)
# 创建画布图片
img_draw = img_origin.copy()
img_draw_former = img_origin.copy()
# 初始化用于计算Fd所需的变量
seedX = -1
seedY = -1
seedX_origin = -1
seedY_origin = -1
maxDp = 0
maxDq = 0
# 展示程序窗口
print("LClick to set seed,RClick to finish,press Q to quit")
processing = False
while True:
    cv2.imshow("wind", img_draw)
    if cv2.waitKey(20) == ord('q'):
        break

# 退出并清理所有窗口
cv2.destroyAllWindows()
